import os
os.environ['CUDA_VISIBLE_DEVICES'] = '1'
os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'
import torch
from PIL import Image
import open_clip
import pandas as pd
import numpy as np
from tqdm import tqdm
import pyarrow as pa
import pyarrow.parquet as pq

model, _, preprocess_val = open_clip.create_model_and_transforms('hf-hub:Marqo/marqo-fashionSigLIP')
model.cuda()
model.eval()

imgs_root = ['/home/tfj/datasets/image_retri10k/eval_images_o/',
             '/home/tfj/datasets/image_retri10k/eval_images_q/',
             '/home/tfj/datasets/image_retri10k/eval_images_v/']
all_images = pd.read_csv('data/merged_table/goods_images_eval_I_old.csv')
images_id = all_images['image_id'].to_list()
goods_id = all_images['goods_id'].to_list()
imgs_url = all_images['image_url'].to_list()
image_names = []
image_paths = []
imgs_urls = []
for img_id, good_id in zip(images_id, goods_id):
    image_name = f'{good_id}_{img_id}.jpg'
    found = False
    for rt in imgs_root:
        img_path = os.path.join(rt, image_name)
        if os.path.isfile(img_path):
            image_names.append(image_name)
            image_paths.append(img_path)
            imgs_urls.append(img_path)
            found = True
            break
    if not found:
        # 可以记录未找到的图片
        pass

batch_size = 256
features = []

with torch.no_grad():
    for i in tqdm(range(0, len(image_paths), batch_size)):
        batch_paths = image_paths[i:i+batch_size]
        batch_imgs = []
        for p in batch_paths:
            img = Image.open(p).convert('RGB')
            batch_imgs.append(preprocess_val(img))
        batch_tensor = torch.stack(batch_imgs).cuda()
        batch_features = model.encode_image(batch_tensor, normalize=True)
        batch_features = batch_features.cpu().numpy()
        features.append(batch_features)
        # break
features = np.concatenate(features, axis=0)

# 保存为parquet
table = pa.Table.from_pydict({
    'image_name': image_names,
    'imgs_url': imgs_urls,
    'features': [f.astype(np.float32) for f in features]
})
pq.write_table(table, 'fashion_clip_features.parquet')


